package com.example.demo.handler;

import com.example.demo.model.ChatMessage;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;
import org.springframework.beans.factory.annotation.Autowired;
import org.springframework.stereotype.Component;
import org.springframework.web.socket.CloseStatus;
import org.springframework.web.socket.TextMessage;
import org.springframework.web.socket.WebSocketSession;
import org.springframework.web.socket.handler.TextWebSocketHandler;
import java.util.Map;
import java.util.concurrent.ConcurrentHashMap;
import com.alibaba.fastjson.JSON;
import com.alibaba.fastjson.JSONObject;
import com.example.demo.service.ChatService;
import java.util.Date;

@Component
public class ChatWebSocketHandler extends TextWebSocketHandler {
    private static final Logger log = LoggerFactory.getLogger(ChatWebSocketHandler.class);
    
    // 存储用户的WebSocket会话
    private static final Map<Integer, WebSocketSession> userSessions = new ConcurrentHashMap<>();

    @Autowired
    private ChatService chatService;

    @Override
    public void afterConnectionEstablished(WebSocketSession session) {
        Integer userId = getUserId(session);
        if (userId != null) {
            userSessions.put(userId, session);
            log.info("用户{}建立WebSocket连接", userId);
        }
    }

    @Override
    protected void handleTextMessage(WebSocketSession session, TextMessage message) throws Exception {
        ChatMessage chatMessage = JSON.parseObject(message.getPayload(), ChatMessage.class);
        Integer senderId = getUserId(session);
        String senderName = (String) session.getAttributes().get("userName");
        
        // 设置发送者信息
        chatMessage.setSenderId(senderId);
        chatMessage.setSenderName(senderName);
        chatMessage.setCreateTime(new Date());
        
        // 保存消息到数据库
        saveChatMessage(chatMessage);
        
        // 发送给接收者
        WebSocketSession receiverSession = userSessions.get(chatMessage.getReceiverId());
        if (receiverSession != null && receiverSession.isOpen()) {
            receiverSession.sendMessage(new TextMessage(JSON.toJSONString(chatMessage)));
        }
    }

    @Override
    public void afterConnectionClosed(WebSocketSession session, CloseStatus status) {
        Integer userId = getUserId(session);
        if (userId != null) {
            userSessions.remove(userId);
            log.info("用户{}断开WebSocket连接", userId);
        }
    }

    private Integer getUserId(WebSocketSession session) {
        Map<String, Object> attributes = session.getAttributes();
        return (Integer) attributes.get("userId");
    }

    private void saveChatMessage(ChatMessage message) {
        chatService.saveMessage(message);
    }
} 